import os
import cv2
import numpy as np
from PIL import Image
from tqdm import tqdm

# 原始图像路径
image_path = "/home/t/workspace/XMem/workspace/177661140-f690156b-1775-4cd7-acd7-1738a5c92f30/images"
# mask 文件夹路径
mask_path = "./logs/mask0/"

def load_masks_and_images(image_path, mask_path):
    image_files = sorted([f for f in os.listdir(image_path) if f.endswith('.jpg')])
    masks_list = []
    bboxes_list = []
    images_list = []

    # print(image_files)

    # 使用 tqdm 进度条
    pbar = tqdm(total=len(image_files), unit='image')

    for image_file in image_files:
        mask_dir = os.path.splitext(image_file)[0]
        mask_dir_path = os.path.join(mask_path, mask_dir)

        # print("image_file:", image_file)
        # print("mask_path:", mask_path)
        # print("mask_dir:", mask_dir)

        if os.path.isdir(mask_dir_path):
            mask_files = sorted([f for f in os.listdir(mask_dir_path) if f.endswith('.png')])
            masks = []
            bboxes = []

            # print(mask_files)

            for mask_file in mask_files:
                mask_path1 = os.path.join(mask_dir_path, mask_file)
                # print("mask_path:", mask_path1)
                mask = np.array(cv2.imread(mask_path1))
                # print("shape:", mask.shape)
                masks.append(mask)

                # 计算边界框
                rows = np.any(mask, axis=1)
                cols = np.any(mask, axis=0)
                rmin, rmax = np.where(rows)[0][[0, -1]]
                cmin, cmax = np.where(cols)[0][[0, -1]]
                bbox = [cmin, rmin, cmax, rmax]
                bboxes.append(bbox)

            masks_list.append(masks)
            bboxes_list.append(bboxes)

            image_path1 = os.path.join(image_path, image_file)
            image = cv2.imread(image_path1)
            images_list.append(image)

        pbar.update(1)

    return masks_list, bboxes_list, images_list

def draw_bboxes(image, bboxes):
    for bbox in bboxes:
        x1, y1, x2, y2 = bbox
        cv2.rectangle(image, (x1, y1), (x2, y2), (0, 0, 255), 2)
    return image

def draw_masks(image, masks):
    for mask in masks:
        image = cv2.addWeighted(image, 1, mask, 0.5, 0)
    return image

def main():
    masks_list, bboxes_list, images_list = load_masks_and_images(image_path, mask_path)

    # 显示 masks_list, bboxes_list, images_list 各自的形状
    print("masks_list shape:", len(masks_list))
    print("bboxes_list shape:", len(bboxes_list))
    print("images_list shape:", len(images_list))

    for masks, bboxes, image in zip(masks_list, bboxes_list, images_list):
        image_with_bboxes = draw_bboxes(image, bboxes)
        image_with_bboxes = draw_masks(image_with_bboxes, masks)
        cv2.imshow("Image with Bounding Boxes", image_with_bboxes)
        cv2.waitKey(0)

if __name__ == "__main__":
    main()